import os
import pickle
import cv2
import numpy as np
from tqdm import tqdm
import utils


def calc_mean_std_func(img_dir, pkl_file="mean_std.pkl"):

    #如果 pkl 存在，直接读取
    script_dir = os.path.dirname(os.path.abspath(__file__))
    pkl_file = os.path.join(script_dir, "mean_std.pkl")
    if os.path.exists(pkl_file):
        with open(pkl_file, "rb") as f:
            data = pickle.load(f)
        mean = data["mean"]
        std = data["std"]
        print(f"已读取缓存 pkl 文件: {pkl_file}")
        return mean, std


    img_paths = []

    for file in tqdm(os.listdir(img_dir), desc="正在获取图像路径"):
        if file.endswith('.png') and "matte" not in file:
            img_paths.append(os.path.join(img_dir, file))

    mean = np.zeros(3) # RGB三通道
    std = np.zeros(3)
    samples_len = 0

    for img_path in tqdm(img_paths, desc="正在计算均值和标准差"):
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR 转 RGB
        img = img.astype(np.float32) / 255.0 # 归一化

        mean += np.mean(img, axis=(0, 1)) # 高度 -> 长度为3的向量
        std += np.std(img, axis=(0, 1)) # 宽度 -> 长度为3的向量

        samples_len += 1

    mean /= samples_len
    std /= samples_len

    mean_list = mean.tolist()
    std_list = std.tolist()

    utils.save_mean_std_to_pkl(mean_list, std_list, pkl_file)

    return mean_list, std_list

if __name__ == "__main__":
    img_dir = "../Portrait-dataset-2000/dataset/training"

    mean, std = calc_mean_std_func(img_dir)
    print("Mean:", mean)
    print("Std:", std)

